{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from pydantic import BaseModel, ValidationError, Field\n",
    "\n",
    "import json \n",
    "\n",
    "from typing import Annotated, List\n",
    "from openai import AsyncOpenAI\n",
    "\n",
    "\n",
    "from semantic_kernel.kernel import Kernel\n",
    "from semantic_kernel.connectors.ai.open_ai import OpenAIChatCompletion, OpenAIChatPromptExecutionSettings\n",
    "from semantic_kernel.agents import ChatCompletionAgent\n",
    "from semantic_kernel.contents import ChatHistory\n",
    "\n",
    "\n",
    "from semantic_kernel.agents.open_ai import OpenAIAssistantAgent\n",
    "from semantic_kernel.contents import AuthorRole, ChatMessageContent\n",
    "from semantic_kernel.functions import kernel_function\n",
    "\n",
    "from semantic_kernel.connectors.ai import FunctionChoiceBehavior\n",
    "\n",
    "from semantic_kernel.contents.function_call_content import FunctionCallContent\n",
    "from semantic_kernel.contents.function_result_content import FunctionResultContent\n",
    "from semantic_kernel.functions import KernelArguments, kernel_function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv()\n",
    "client = AsyncOpenAI(\n",
    "    api_key=os.environ.get(\"GITHUB_TOKEN\"), base_url=\"https://models.inference.ai.azure.com/\")\n",
    "\n",
    "kernel = Kernel()\n",
    "\n",
    "\n",
    "chat_completion_service = OpenAIChatCompletion(\n",
    "    ai_model_id=\"gpt-4o-mini\",\n",
    "    async_client=client,\n",
    "  \n",
    ")\n",
    "kernel.add_service(chat_completion_service)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SubTask(BaseModel):\n",
    "    assigned_agent: str = Field(\n",
    "        description=\"The specific agent assigned to handle this subtask\")\n",
    "    task_details: str = Field(\n",
    "        description=\"Detailed description of what needs to be done for this subtask\")\n",
    "\n",
    "\n",
    "class TravelPlan(BaseModel):\n",
    "    main_task: str = Field(\n",
    "        description=\"The overall travel request from the user\")\n",
    "    subtasks: List[SubTask] = Field(\n",
    "        description=\"List of subtasks broken down from the main task, each assigned to a specialized agent\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "AGENT_NAME = \"TravelAgent\"\n",
    "AGENT_INSTRUCTIONS = \"\"\"You are an planner agent.\n",
    "    Your job is to decide which agents to run based on the user's request.\n",
    "    Below are the available agents specialised in different tasks:\n",
    "    - FlightBooking: For booking flights and providing flight information\n",
    "    - HotelBooking: For booking hotels and providing hotel information\n",
    "    - CarRental: For booking cars and providing car rental information\n",
    "    - ActivitiesBooking: For booking activities and providing activity information\n",
    "    - DestinationInfo: For providing information about destinations\n",
    "    - DefaultAgent: For handling general requests\"\"\"\n",
    "\n",
    "# Get prompt execution settings and configure response format\n",
    "settings = kernel.get_prompt_execution_settings_from_service_id(\"default\")\n",
    "assert isinstance(settings, OpenAIChatPromptExecutionSettings)\n",
    "settings.response_format = TravelPlan\n",
    "\n",
    "agent = ChatCompletionAgent(\n",
    "    kernel=kernel,\n",
    "    name=AGENT_NAME,\n",
    "    instructions=AGENT_INSTRUCTIONS,\n",
    "    arguments=KernelArguments(settings) \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div style='margin-bottom:10px'><div style='font-weight:bold'>User:</div><div style='margin-left:20px'>Create a travel plan for a family of 2 kids from Singapore to Melboune</div></div><div style='margin-bottom:20px'><div style='font-weight:bold'>Validated Travel Plan:</div><pre style='margin-left:20px; padding:10px; border-radius:5px;'>{\n",
       "  \"main_task\": \"Travel plan for a family trip from Singapore to Melbourne\",\n",
       "  \"subtasks\": [\n",
       "    {\n",
       "      \"assigned_agent\": \"FlightBooking\",\n",
       "      \"task_details\": \"Book round-trip flights from Singapore to Melbourne for a family of four (2 adults, 2 kids) including preferred flight times and airlines.\"\n",
       "    },\n",
       "    {\n",
       "      \"assigned_agent\": \"HotelBooking\",\n",
       "      \"task_details\": \"Find and book family-friendly accommodation in Melbourne for four guests, ideally near local attractions and amenities.\"\n",
       "    },\n",
       "    {\n",
       "      \"assigned_agent\": \"CarRental\",\n",
       "      \"task_details\": \"Arrange for a rental car in Melbourne for family use during the stay, including options for child car seats.\"\n",
       "    },\n",
       "    {\n",
       "      \"assigned_agent\": \"ActivitiesBooking\",\n",
       "      \"task_details\": \"Plan and book family-friendly activities and attractions in Melbourne suitable for kids, such as zoos, parks, and museums.\"\n",
       "    },\n",
       "    {\n",
       "      \"assigned_agent\": \"DestinationInfo\",\n",
       "      \"task_details\": \"Provide detailed information about Melbourne, including weather, local cuisine, and family-friendly tips.\"\n",
       "    }\n",
       "  ]\n",
       "}</pre></div><hr>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import display, HTML\n",
    "\n",
    "\n",
    "async def main():\n",
    "    # Define the chat history\n",
    "    chat_history = ChatHistory()\n",
    "\n",
    "    # Respond to user input\n",
    "    user_inputs = [\n",
    "        \"Create a travel plan for a family of 2 kids from Singapore to Melboune\",\n",
    "    ]\n",
    "\n",
    "    for user_input in user_inputs:\n",
    "        # Add the user input to the chat history\n",
    "        chat_history.add_user_message(user_input)\n",
    "\n",
    "        # Start building HTML output\n",
    "        html_output = f\"<div style='margin-bottom:10px'>\"\n",
    "        html_output += f\"<div style='font-weight:bold'>User:</div>\"\n",
    "        html_output += f\"<div style='margin-left:20px'>{user_input}</div>\"\n",
    "        html_output += f\"</div>\"\n",
    "\n",
    "        agent_name: str | None = None\n",
    "        full_response = \"\"\n",
    "\n",
    "        # Collect the agent's response\n",
    "        async for content in agent.invoke_stream(chat_history):\n",
    "            if not agent_name and hasattr(content, 'name'):\n",
    "                agent_name = content.name\n",
    "            # Add this line to collect the response content\n",
    "            if hasattr(content, 'content') and content.content is not None:\n",
    "                full_response += content.content\n",
    "\n",
    "        try:\n",
    "            # Try to validate the response as a TravelPlan\n",
    "            travel_plan = TravelPlan.model_validate_json(full_response)\n",
    "\n",
    "            # Display the validated model as formatted JSON\n",
    "            formatted_json = json.dumps(travel_plan.model_dump(), indent=2)\n",
    "            html_output += f\"<div style='margin-bottom:20px'>\"\n",
    "            html_output += f\"<div style='font-weight:bold'>Validated Travel Plan:</div>\"\n",
    "            html_output += f\"<pre style='margin-left:20px; padding:10px; border-radius:5px;'>{formatted_json}</pre>\"\n",
    "            html_output += f\"</div>\"\n",
    "        except ValidationError as e:\n",
    "            # Handle validation errors\n",
    "            html_output += f\"<div style='margin-bottom:20px; color:red;'>\"\n",
    "            html_output += f\"<div style='font-weight:bold'>Validation Error:</div>\"\n",
    "            html_output += f\"<pre style='margin-left:20px;'>{str(e)}</pre>\"\n",
    "            html_output += f\"</div>\"\n",
    "            # Add this to see what the response contains for debugging\n",
    "            html_output += f\"<div style='margin-bottom:20px;'>\"\n",
    "            html_output += f\"<div style='font-weight:bold'>Raw Response:</div>\"\n",
    "            html_output += f\"<div style='margin-left:20px; white-space:pre-wrap'>{full_response}</div>\"\n",
    "            html_output += f\"</div>\"\n",
    "\n",
    "        html_output += \"<hr>\"\n",
    "\n",
    "        # Display formatted HTML\n",
    "        display(HTML(html_output))\n",
    "\n",
    "await main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
